"""
模型分发路由
实现feed端调用pyolo端下载和分发模型的功能
"""

from fastapi import APIRouter, HTTPException, BackgroundTasks
from typing import List, Optional, Dict, Any
import time
import aiohttp
import asyncio
from pydantic import BaseModel

from ..storage import data_manager
from ..logger import logger
from ..services.event_dispatcher import event_dispatcher, EventType
from ..utils import generate_id

router = APIRouter()


class ModelDispatchRequest(BaseModel):
    """模型分发请求"""
    model_name: str
    model_url: str
    model_type: str = "detection"
    version: str = "1.0.0"
    description: str = ""
    author: str = ""
    target_clients: Optional[List[str]] = None  # 指定目标客户端，None表示所有客户端


class ModelDispatchResponse(BaseModel):
    """模型分发响应"""
    success: bool
    message: str
    dispatch_id: str
    target_clients: List[str]
    model_info: Dict[str, Any]


class DispatchStatus(BaseModel):
    """分发状态"""
    dispatch_id: str
    model_name: str
    status: str  # pending, downloading, completed, failed
    progress: float  # 0-100
    target_clients: List[str]
    completed_clients: List[str]
    failed_clients: List[str]
    start_time: float
    end_time: Optional[float] = None
    error_message: Optional[str] = None


# 存储分发状态
dispatch_statuses: Dict[str, DispatchStatus] = {}


async def call_pyolo_download_model(
    client_info: Dict[str, Any], 
    model_name: str, 
    model_url: str,
    model_type: str,
    version: str,
    description: str,
    author: str
) -> Dict[str, Any]:
    """调用pyolo客户端下载模型"""
    try:
        # 构建pyolo客户端URL
        base_url = f"http://{client_info['ip_address']}:{client_info['port']}"
        
        # 准备请求数据
        form_data = {
            "url": model_url,
            "name": model_name,
            "model_type": model_type,
            "version": version,
            "description": description,
            "author": author
        }
        
        # 调用pyolo的模型上传接口
        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{base_url}/models/upload-url",
                data=form_data,
                timeout=aiohttp.ClientTimeout(total=300)  # 5分钟超时
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    return {
                        "success": True,
                        "client_id": client_info["client_id"],
                        "result": result
                    }
                else:
                    error_text = await response.text()
                    return {
                        "success": False,
                        "client_id": client_info["client_id"],
                        "error": f"HTTP {response.status}: {error_text}"
                    }
                    
    except asyncio.TimeoutError:
        return {
            "success": False,
            "client_id": client_info["client_id"],
            "error": "Request timeout"
        }
    except Exception as e:
        return {
            "success": False,
            "client_id": client_info["client_id"],
            "error": str(e)
        }


async def dispatch_model_to_clients(
    dispatch_id: str,
    model_name: str,
    model_url: str,
    model_type: str,
    version: str,
    description: str,
    author: str,
    target_clients: List[str]
):
    """分发模型到指定客户端"""
    try:
        # 更新分发状态
        dispatch_statuses[dispatch_id].status = "downloading"
        dispatch_statuses[dispatch_id].progress = 0
        
        # 获取目标客户端信息
        from .clients import get_pyolo_client_info
        all_clients = await get_pyolo_client_info()
        
        if target_clients:
            # 筛选指定客户端
            target_client_infos = [
                client for client in all_clients 
                if client["client_id"] in target_clients
            ]
        else:
            # 所有在线客户端
            target_client_infos = all_clients
        
        if not target_client_infos:
            dispatch_statuses[dispatch_id].status = "failed"
            dispatch_statuses[dispatch_id].error_message = "No target clients found"
            return
        
        # 并发调用所有客户端
        tasks = []
        for client_info in target_client_infos:
            task = call_pyolo_download_model(
                client_info, model_name, model_url, model_type, 
                version, description, author
            )
            tasks.append(task)
        
        # 等待所有任务完成
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果
        completed_clients = []
        failed_clients = []
        
        for i, result in enumerate(results):
            client_id = target_client_infos[i]["client_id"]
            
            if isinstance(result, Exception):
                failed_clients.append({
                    "client_id": client_id,
                    "error": str(result)
                })
            elif result.get("success"):
                completed_clients.append(client_id)
            else:
                failed_clients.append({
                    "client_id": client_id,
                    "error": result.get("error", "Unknown error")
                })
        
        # 更新分发状态
        dispatch_statuses[dispatch_id].completed_clients = completed_clients
        dispatch_statuses[dispatch_id].failed_clients = [f["client_id"] for f in failed_clients]
        dispatch_statuses[dispatch_id].progress = 100
        
        if failed_clients:
            dispatch_statuses[dispatch_id].status = "completed_with_errors"
            dispatch_statuses[dispatch_id].error_message = f"Failed on {len(failed_clients)} clients"
        else:
            dispatch_statuses[dispatch_id].status = "completed"
        
        dispatch_statuses[dispatch_id].end_time = time.time()
        
        # 发出分发完成事件
        await event_dispatcher.emit(
            EventType.MODEL_DISPATCHED,
            {
                "dispatch_id": dispatch_id,
                "model_name": model_name,
                "completed_clients": completed_clients,
                "failed_clients": failed_clients,
                "status": dispatch_statuses[dispatch_id].status
            },
            "model_dispatch"
        )
        
        logger.info(f"Model dispatch {dispatch_id} completed: {len(completed_clients)} success, {len(failed_clients)} failed")
        
    except Exception as e:
        logger.error(f"Model dispatch {dispatch_id} failed: {e}")
        dispatch_statuses[dispatch_id].status = "failed"
        dispatch_statuses[dispatch_id].error_message = str(e)
        dispatch_statuses[dispatch_id].end_time = time.time()


@router.post("/dispatch", response_model=ModelDispatchResponse)
async def dispatch_model(
    request: ModelDispatchRequest,
    background_tasks: BackgroundTasks
):
    """分发模型到pyolo客户端"""
    try:
        # 生成分发ID
        dispatch_id = generate_id()
        
        # 获取目标客户端
        if request.target_clients:
            target_clients = request.target_clients
        else:
            # 获取所有在线pyolo客户端
            from .clients import get_pyolo_client_info
            all_clients = await get_pyolo_client_info()
            target_clients = [client["client_id"] for client in all_clients]
        
        if not target_clients:
            raise HTTPException(
                status_code=400, 
                detail="No target clients available for model dispatch"
            )
        
        # 创建分发状态记录
        dispatch_statuses[dispatch_id] = DispatchStatus(
            dispatch_id=dispatch_id,
            model_name=request.model_name,
            status="pending",
            progress=0,
            target_clients=target_clients,
            completed_clients=[],
            failed_clients=[],
            start_time=time.time()
        )
        
        # 在后台执行分发任务
        background_tasks.add_task(
            dispatch_model_to_clients,
            dispatch_id,
            request.model_name,
            request.model_url,
            request.model_type,
            request.version,
            request.description,
            request.author,
            target_clients
        )
        
        # 创建模型记录到feed存储
        from ..storage import ModelInfo
        model_info = ModelInfo(
            id=generate_id(),
            name=request.model_name,
            version=request.version,
            model_type=request.model_type,
            description=request.description,
            author=request.author,
            size="m",  # 默认中等大小
            upload_time=time.time(),
            last_update=time.time(),
            status="dispatching",
            download_url=request.model_url,
            checksum="",
            server_id="dispatch",
            tags=[],
            requirements={},
            metrics={}
        )
        
        await data_manager.add_model(model_info)
        
        # 发出模型分发事件
        await event_dispatcher.emit(
            EventType.MODEL_DISPATCH_STARTED,
            {
                "dispatch_id": dispatch_id,
                "model_name": request.model_name,
                "target_clients": target_clients
            },
            "model_dispatch"
        )
        
        logger.info(f"Started model dispatch {dispatch_id} for {request.model_name} to {len(target_clients)} clients")
        
        return ModelDispatchResponse(
            success=True,
            message=f"Model dispatch started for {request.model_name}",
            dispatch_id=dispatch_id,
            target_clients=target_clients,
            model_info={
                "name": request.model_name,
                "type": request.model_type,
                "version": request.version,
                "url": request.model_url
            }
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Failed to start model dispatch: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to start model dispatch: {str(e)}")


@router.get("/dispatch/{dispatch_id}/status")
async def get_dispatch_status(dispatch_id: str):
    """获取分发状态"""
    if dispatch_id not in dispatch_statuses:
        raise HTTPException(status_code=404, detail="Dispatch not found")
    
    return dispatch_statuses[dispatch_id]


@router.get("/dispatch")
async def get_all_dispatch_statuses():
    """获取所有分发状态"""
    return {
        "dispatches": list(dispatch_statuses.values()),
        "total": len(dispatch_statuses)
    }


@router.post("/dispatch/{dispatch_id}/cancel")
async def cancel_dispatch(dispatch_id: str):
    """取消分发任务"""
    if dispatch_id not in dispatch_statuses:
        raise HTTPException(status_code=404, detail="Dispatch not found")
    
    dispatch = dispatch_statuses[dispatch_id]
    
    if dispatch.status in ["completed", "failed", "cancelled"]:
        raise HTTPException(
            status_code=400, 
            detail=f"Cannot cancel dispatch in status: {dispatch.status}"
        )
    
    dispatch.status = "cancelled"
    dispatch.end_time = time.time()
    
    logger.info(f"Cancelled model dispatch {dispatch_id}")
    
    return {
        "success": True,
        "message": f"Dispatch {dispatch_id} cancelled"
    }

